from typing import Union, List, Tuple

import torch
import warnings
from torch_geometric.data import Data, InMemoryDataset
# 进度条
from tqdm import tqdm
from sklearn.preprocessing import LabelEncoder
import pandas as pd
import numpy as np
from torch_geometric.nn import TopKPooling, SAGEConv
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
import torch.nn.functional as F
from torch_geometric.loader import DataLoader
from sklearn.metrics import roc_auc_score

warnings.filterwarnings("ignore")


def simple():
    x = torch.tensor([[2, 1], [5, 6], [3, 7], [12, 0]], dtype=torch.float)
    y = torch.tensor([0, 1, 0, 1], dtype=torch.float)
    # 边的属性
    edge_index = torch.tensor([[0, 1, 2, 0, 3]  # 起始点
                                  , [1, 0, 1, 3, 2]], dtype=torch.long)  # 终止 点

    data = Data(x=x, y=y, edge_index=edge_index)
    print(data)


def example():
    '''
    yoochoose-clicks.dat：表示用户的浏览行为，其中一个session_id就表示一次登录都浏览了啥东西。
    yoochoose-buys.dat:表示用户最终购买了啥，里面是标签。
    '''
    df = pd.read_csv('../data/yoochoose-click.dat', header=None)
    df.columns = ['session_id', 'timestamp', 'item_id', 'category']

    buy_df = pd.read_csv('../data/yoochoose-buys.dat', header=None)
    buy_df.columns = ['session_id', 'timestamp', 'item_id', 'price', 'quantity']

    item_encoder = LabelEncoder()
    df['item_id'] = item_encoder.fit_transform(df.item_id)
    print(df.head())
    # 随机拿部分数据去做
    sampled_session_id = np.random.choice(df.session_id.unique(), 500, replace=False)
    df = df.loc[df.session_id.isin(sampled_session_id)]
    print(df.nunique())

    # 获取标签
    df['label'] = df.session_id.isin(buy_df.session_id)
    print(df.head())

    # 制作数据集的格式
    data_list = []
    df_test = df[:100]
    grouped = df_test.groupby('session_id')
    for session_id, group in tqdm(grouped):
        print('session_id:', session_id)
        sess_item_id = LabelEncoder().fit_transform(group.item_id)
        print('sess_item_id:', sess_item_id)
        group = group.reset_index(drop=True)
        group['sess_item_id'] = sess_item_id
        print('group:', group)
        node_features = group.loc[group.session_id == session_id, ['sess_item_id', 'item_id']].sort_values(
            'sess_item_id').item_id.drop_duplicates().values
        node_features = torch.LongTensor(node_features).unsqueeze(1)
        print('node_features', node_features)
        target_nodes = group.sess_item_id.values[1:]
        source_nodes = group.sess_item_id.values[:-1]
        print('target_nodes:', target_nodes)
        print('source_nodes:', source_nodes)
        edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)
        x = node_features
        y = torch.FloatTensor([group.label.values[0]])
        data = Data(x=x, y=y, edge_index=edge_index)
        print('data:', data)
        data_list.append(data)


df = pd.read_csv('../data/yoochoose-click.dat', header=None)
df.columns = ['session_id', 'timestamp', 'item_id', 'category']
buy_df = pd.read_csv('../data/yoochoose-buys.dat', header=None)
buy_df.columns = ['session_id', 'timestamp', 'item_id', 'price', 'quantity']

item_encoder = LabelEncoder()
df['item_id'] = item_encoder.fit_transform(df.item_id)
# print(df.head())
# 随机拿部分数据去做
# sampled_session_id = np.random.choice(df.session_id.unique(), 500, replace=False)
# df = df.loc[df.session_id.isin(sampled_session_id)]
# print(df.nunique())

# 获取标签
df['label'] = df.session_id.isin(buy_df.session_id)


class YooChooseBinaryDataSet(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(YooChooseBinaryDataSet, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return []

    @property
    def processed_file_names(self):
        return ['yoochoose_click_binary_1M_sess.dataset']

    def download(self) -> None:
        pass

    def process(self) -> None:
        data_list = []
        grouped = df.groupby('session_id')
        # 遍历每一组session_id（每一个图），目的是将其制作成(from torch_geometric.data import Data)格式
        for session_id, group in tqdm(grouped):
            # 将item_id转换成sess_item_id，从0开始。对每一组session_id中的所有item_id进行编码(例如15453,3651,15452)就按照数值大小编码成(2,0,1)，索引是从0开始的
            sess_item_id = LabelEncoder().fit_transform(group.item_id)
            group = group.reset_index(drop=True)
            group['sess_item_id'] = sess_item_id
            node_features = group.loc[group.session_id == session_id, ['sess_item_id', 'item_id']].sort_values(
                'sess_item_id').item_id.drop_duplicates().values  # item_id作为node_features

            node_features = torch.LongTensor(node_features).unsqueeze(1)
            target_nodes = group.sess_item_id.values[1:]
            # 构建邻接矩阵，第二个到最后一个（点的特征就由其ID组成，edge_index是这样，因为咱们浏览的过程中是有顺序的比如(0,0,2,1)- 5.
            # 所以边就是0->0, 0->2, 2->1这样的，对应的索引就为target_nodes: [0 2 1]，source_nodes: [0 0 2]）
            source_nodes = group.sess_item_id.values[:-1]  # 第一个到倒数第二个

            edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)  # 传入形成二维邻接矩阵
            x = node_features
            y = torch.FloatTensor([group.label.values[0]])
            data = Data(x=x, edge_index=edge_index, y=y)  # 传入参数，构建data
            data_list.append(data)  # 每一个图都放入data_list

        data, slices = self.collate(data_list)  # 转换成可以保存到本地的格式
        torch.save((data, slices), self.processed_paths[0])  # 保存数据


class Net(torch.nn.Module):  # 针对图进行分类任务
    def __init__(self, embed_dim=128):
        super(Net, self).__init__()

        self.conv1 = SAGEConv(embed_dim, 128)  # 卷积层 输入embed_dim，输出128
        self.pool1 = TopKPooling(128, ratio=0.8)  # 做剪枝操作
        self.conv2 = SAGEConv(128, 128)
        self.pool2 = TopKPooling(128, ratio=0.8)
        self.conv3 = SAGEConv(128, 128)
        self.pool3 = TopKPooling(128, ratio=0.8)
        self.item_embedding = torch.nn.Embedding(num_embeddings=df.item_id.max() + 10, embedding_dim=embed_dim)  # 映射向量

        self.lin1 = torch.nn.Linear(128, 128)
        self.lin2 = torch.nn.Linear(128, 64)
        self.lin3 = torch.nn.Linear(64, 1)

        self.bn1 = torch.nn.BatchNorm1d(128)
        self.bn2 = torch.nn.BatchNorm1d(64)
        self.act1 = torch.nn.ReLU()
        self.act2 = torch.nn.ReLU()

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch  # x:n*1,其中每个图里点的个数是不同的
        # print(x)
        x = self.item_embedding(x)  # n*1*128 特征编码后的结果
        # print('item_embedding',x.shape)
        x = x.squeeze(1)  # n*128
        # print('squeeze',x.shape)

        # 对输入不断做卷积，不断做池化池化，得到的特征会越来越浓缩，图会越来越小，但是池化完成之后的特征维度都是一样的
        x = F.relu(self.conv1(x, edge_index))  # n*128
        # print('conv1',x.shape)
        x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)  # pool之后得到 n*0.8个点
        # print('self.pool1',x.shape)
        # print('self.pool1',edge_index)
        # print('self.pool1',batch)
        # x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x1 = gap(x, batch)  # gap:全局平均池化  得到全局特征
        # print('gmp',gmp(x, batch).shape) # batch*128
        # print('cat',x1.shape) # batch*256
        x = F.relu(self.conv2(x, edge_index))
        # print('conv2',x.shape)
        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
        # print('pool2',x.shape)
        # print('pool2',edge_index)
        # print('pool2',batch)
        # x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x2 = gap(x, batch)
        # print('x2',x2.shape)
        x = F.relu(self.conv3(x, edge_index))
        # print('conv3',x.shape)
        x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)
        # print('pool3',x.shape)
        # x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x3 = gap(x, batch)
        # print('x3',x3.shape)# batch * 256
        x = x1 + x2 + x3  # 获取不同尺度的全局特征
        """通过全连接层，得到最终输出结果值"""
        x = self.lin1(x)
        # print('lin1',x.shape)
        x = self.act1(x)
        x = self.lin2(x)
        # print('lin2',x.shape)
        x = self.act2(x)
        x = F.dropout(x, p=0.5, training=self.training)

        x = torch.sigmoid(self.lin3(x)).squeeze(1)  # batch个结果
        # print('sigmoid',x.shape)
        return x


def train(dataset, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    crit = torch.nn.BCELoss()
    train_loader = DataLoader(dataset, batch_size=64)

    for epoch in range(10):
        model.train()
        loss_all = 0
        for data in train_loader:  # 遍历dataloader
            # print('data',data)
            optimizer.zero_grad()
            output = model(data)  # data数据传入模型
            label = data.y
            loss = crit(output, label)  # 计算损失
            loss.backward()
            loss_all += data.num_graphs * loss.item()
            optimizer.step()  # 梯度更新
        loss = loss_all / len(dataset)
        print('loss:', loss)


def evalute(loader, model):
    for epoch in range(1):
        model.eval()
        prediction = []
        labels = []
        with torch.no_grad():
            for data in loader:
                data = data  # .to(device)
                pred = model(data)  # .detach().cpu().numpy()

                label = data.y  # .detach().cpu().numpy()
                prediction.append(pred)
                labels.append(label)
        prediction = np.hstack(prediction)
        labels = np.hstack(labels)
        score = roc_auc_score(labels, prediction)
        print('roc_auc_score:', score)


if __name__ == '__main__':
    # simple()
    # example()

    dataset = YooChooseBinaryDataSet(root='../data/')
    model = Net()
    train(dataset, model)
    evalute(dataset, model)
